#include "yolo_v6_o.h"
#include "yaml-cpp/yaml.h"
#include "boost/process.hpp"
#include "common/device_utility.h"
#include "boost/filesystem/operations.hpp"
#include "boost/filesystem/path.hpp"
#include "boost/filesystem/fstream.hpp"

using namespace CryptoPP;

std::vector<char> decryptENNX(
    const std::vector<char> & v_encrypt,
    const std::string& key,
    const std::string& iv) {
    std::vector<char> v_decrypt;
    try {
        CBC_Mode<AES>::Decryption decryption((const byte*)key.data(), key.size(), (const byte*)iv.data());

        std::vector<char> _decrypt = v_encrypt;
        decryption.ProcessData(
            reinterpret_cast<CryptoPP::byte*>(_decrypt.data()),
            reinterpret_cast<const CryptoPP::byte*>(v_encrypt.data()),
            v_encrypt.size());

        byte paddingByte = _decrypt.back();
        size_t paddingSize = static_cast<size_t>(paddingByte);
        _decrypt.resize(v_encrypt.size() - paddingSize);

        v_decrypt = _decrypt;
    } catch (const CryptoPP::Exception& e) {
        std::cerr << e.what() << std::endl;
    }
    return v_decrypt;
}

/**
 * @brief Read file data into a character vector
 * @param file_path File path
 * @return A vector of characters containing the file's content
 */
std::vector<char> get_file_data(const std::string& file_path) {
    std::ifstream file(file_path, std::ios_base::binary);
    if (!file.is_open()) {
        throw std::runtime_error("Failed to open file: " + file_path);
    }
    file.seekg(0, std::ios_base::end);
    const size_t file_len = file.tellg();
    file.seekg(0, std::ios_base::beg);
    std::vector<char> data(file_len, 0);
    if (file.read(data.data(), file_len)) {
        return data;
    } else {
        throw std::runtime_error("Failed to read file: " + file_path);
    }
}

/**
 * @brief Write character vector data to a file
 * @param data Data to be written
 * @param file_path File path
 */
void write_file_data(std::vector<char> data, std::string file_path) {
    std::ofstream file(file_path, std::ios_base::binary);
    file.write(data.data(), data.size());
    file.close();
}

void encrypt_data(std::vector<char>& data) {
    char key = 0x1;
    for (size_t i = 0; i < data.size(); ++i) {
        data[i] ^= key;
    }
}

void decrypt_data(std::vector<char>& data) {
    char key = 0x1;
    for (size_t i = 0; i < data.size(); ++i) {
        data[i] ^= key;
    }
}

void merge_files(
    const std::string& xml_path,
    const std::string& bin_path,
    const st_encrypt_info_ovino& info,
    const std::string& output_path) {
    std::vector<char> xml_data = get_file_data(xml_path);
    std::vector<char> bin_data = get_file_data(bin_path);

    encrypt_data(xml_data);
    encrypt_data(bin_data);

    st_encrypt_header header = {
        sizeof(st_encrypt_info_ovino),
        xml_data.size(),
        bin_data.size()
    };

    std::vector<char> merged_data;
    merged_data.reserve(sizeof(st_encrypt_header) + sizeof(st_encrypt_info_ovino) + xml_data.size() + bin_data.size());
    merged_data.insert(merged_data.end(), reinterpret_cast<const char*>(&header), reinterpret_cast<const char*>(&header) + sizeof(st_encrypt_header));
    merged_data.insert(merged_data.end(), reinterpret_cast<const char*>(&info), reinterpret_cast<const char*>(&info) + sizeof(st_encrypt_info_ovino));
    merged_data.insert(merged_data.end(), xml_data.begin(), xml_data.end());
    merged_data.insert(merged_data.end(), bin_data.begin(), bin_data.end());

    std::ofstream output_file(output_path, std::ios_base::binary);
    if (!output_file.is_open()) {
        throw std::runtime_error("Failed to open output file: " + output_path);
    }
    output_file.write(merged_data.data(), merged_data.size());
    if (!output_file) {
        throw std::runtime_error("Failed to write to output file: " + output_path);
    }
}

/**
 * @brief Reads data from a YAML configuration file into an st_encrypt_info_ovino structure
 * @param input Transformation input structure containing the configuration file path and other information
 * @return Populated st_encrypt_info_ovino structure
 * @throws std::runtime_error If there is an issue with the configuration file or if reading fails
 */
st_encrypt_info_ovino loadConfigToInfo(const st_trans_model_input& input) {
    YAML::Node root = YAML::LoadFile(input.cfg_path);
    YAML::Node config;

    if (root["YOLOv6"].IsDefined()) {
        config = root["YOLOv6"];
    } else if (root["DET_LB"].IsDefined()) {
        config = root["DET_LB"];
    } else {
        throw std::runtime_error("Invalid configuration: missing YOLOv6 or DET_LB node.");
    }

    st_encrypt_info_ovino info;

    if (!config["BATCH_SIZE"].IsDefined()) {
        throw std::runtime_error("BATCH_SIZE undefined.");
    }
    info.batch_size = config["BATCH_SIZE"].as<int>();

    if (!config["CLASS_NUM"].IsDefined()) {
        throw std::runtime_error("CLASS_NUM undefined.");
    }
    info.class_num = config["CLASS_NUM"].as<int>();

    if (!config["INPUT_CHANNEL"].IsDefined() ||
        !config["IMAGE_WIDTH"].IsDefined() ||
        !config["IMAGE_HEIGHT"].IsDefined()) {
        throw std::runtime_error("Input channel or image size undefined.");
    }
    info.input_channel = config["INPUT_CHANNEL"].as<int>();
    info.image_width = config["IMAGE_WIDTH"].as<int>();
    info.image_height = config["IMAGE_HEIGHT"].as<int>();

    if (!config["obj_threshold"].IsDefined() ||
        !config["nms_threshold"].IsDefined()) {
        throw std::runtime_error("Detection thresholds undefined.");
    }
    info.obj_threshold = config["obj_threshold"].as<float>();
    info.nms_threshold = config["nms_threshold"].as<float>();

    if (!config["num_anchors"].IsDefined() ||
        !config["strides"].IsDefined()) {
        throw std::runtime_error("Anchors or strides undefined.");
    }

    info.num_anchors_size = root["DET_LB"]["num_anchors"].size();
    for (size_t i = 0; i < info.num_anchors_size; ++i) {
        info.num_anchors[i] = root["DET_LB"]["num_anchors"][i].as<int>();
    }
    info.stride_num = root["DET_LB"]["strides"].size();
    for (size_t i = 0; i < info.stride_num; ++i) {
        info.strides[i] = root["DET_LB"]["strides"][i].as<int>();
    }

    return info;
}

/**
 * @brief Reads a model file and decrypts it (supports .onnx and .ennx files)
 * @param input Transformation input structure containing the model path and other information
 * @param key Decryption key
 * @param iv Decryption initialization vector
 * @return Model file data as a byte vector
 * @throws std::runtime_error If the file format is not supported or if reading fails
 */
std::vector<char> loadAndDecrypt_ONNX_ENNX_Model(
    const st_trans_model_input& input,
    const std::string& key,
    const std::string& iv) {
    namespace fs = boost::filesystem;

    fs::path p(input.model_path.c_str());
    if (p.extension().string() == ".onnx") {
        return get_file_data(p.string());
    } else if (p.extension().string() == ".ennx") {
        std::vector<char> v_e_model = get_file_data(p.string());
        return decryptENNX(v_e_model, key, iv);
    } else {
        throw std::runtime_error("Unsupported model file format: " + p.extension().string());
    }
}

/**
 * @brief Removes temporary model files
 * @param directory The directory to remove temporary model files from (default is the current directory)
 */
void remove_temp_model_files(const std::string& directory = ".") {
    for (const auto& entry : std::filesystem::directory_iterator(directory)) {
        if (entry.is_regular_file()) {
            std::string filename = entry.path().filename().string();
            if (filename.rfind("temp_model.", 0) == 0) {
                try {
                    std::filesystem::remove(entry.path());
                } catch (const std::filesystem::filesystem_error& e) {
                    std::cerr << "Error deleting " << entry.path() << ": " << e.what() << std::endl;
                }
            }
        }
    }
}

/**
 * @brief Converts ONNX model data to IR format XML and BIN file data buffers
 * @param core OpenVINO core object
 * @param model_data ONNX model data as a string
 * @return A pair of buffers containing XML and BIN file data
 */
std::pair<std::vector<char>, std::vector<char>> coreModel_to_XB_buffer(
    ov::Core& core,
    const std::string& model_data) {
    std::shared_ptr<ov::Model> model = core.read_model(model_data, ov::Tensor{});
    std::cout << "Compiled model loaded successfully!" << std::endl;

    const std::string temp_xml_path = "temp_model.xml";
    const std::string temp_bin_path = "temp_model.bin";

    ov::serialize(model, temp_xml_path, temp_bin_path);

    std::ifstream xml_file(temp_xml_path, std::ios::binary);
    std::ifstream bin_file(temp_bin_path, std::ios::binary);

    std::vector<char> xml_buffer((std::istreambuf_iterator<char>(xml_file)), std::istreambuf_iterator<char>());
    std::vector<char> bin_buffer((std::istreambuf_iterator<char>(bin_file)), std::istreambuf_iterator<char>());

    std::remove(temp_xml_path.c_str());
    std::remove(temp_bin_path.c_str());

    return {xml_buffer, bin_buffer};
}

/**
 * @brief Encrypts XML and BIN data and writes them to an .hdats file
 * @param xml_data The XML data buffer to be encrypted
 * @param bin_data The BIN data buffer to be encrypted
 * @param info The encryption information structure
 * @param output_path The path where the encrypted file will be written
 */
void encrypt_XML_BIN_to_hdats(
    std::vector<char>& xml_data,
    std::vector<char>& bin_data,
    const st_encrypt_info_ovino& info,
    const std::string& output_path) {
    encrypt_data(xml_data);
    encrypt_data(bin_data);

    st_encrypt_header header = {
        sizeof(st_encrypt_info_ovino),
        xml_data.size(),
        bin_data.size()
    };

    std::vector<char> merged_data;
    merged_data.reserve(sizeof(st_encrypt_header) + sizeof(st_encrypt_info_ovino) + xml_data.size() + bin_data.size());
    merged_data.insert(merged_data.end(), reinterpret_cast<const char*>(&header), reinterpret_cast<const char*>(&header) + sizeof(st_encrypt_header));
    merged_data.insert(merged_data.end(), reinterpret_cast<const char*>(&info), reinterpret_cast<const char*>(&info) + sizeof(st_encrypt_info_ovino));
    merged_data.insert(merged_data.end(), xml_data.begin(), xml_data.end());
    merged_data.insert(merged_data.end(), bin_data.begin(), bin_data.end());

    std::ofstream output_file(output_path, std::ios_base::binary);
    if (!output_file.is_open()) {
        throw std::runtime_error("Failed to open output file: " + output_path);
    }
    output_file.write(merged_data.data(), merged_data.size());
    if (!output_file) {
        throw std::runtime_error("Failed to write to output file: " + output_path);
    }
}

/**
 * @brief Constructor for the YOLO v6 OpenVINO class
 */
yolo_v6_o::yolo_v6_o() = default;

/**
 * @brief Destructor for the YOLO v6 OpenVINO class
 */
yolo_v6_o::~yolo_v6_o() = default;

/**
 * @brief Converts a YOLO v6 model file using OpenVINO
 * @param input Transformation input structure containing the configuration file path and other information
 * @return true if the conversion is successful, otherwise false
 */
bool yolo_v6_o::TransModel(const st_trans_model_input& input) {
    st_encrypt_info_ovino info;
    std::vector<char> v_model;

    try {
        info = loadConfigToInfo(input);
    } catch (const std::runtime_error& e) {
        std::cerr << "Error loading configuration: " << e.what() << std::endl;
        return false;
    }
    try {
        const std::string key = "irkdh_573%3?iq5h";
        const std::string iv = "1234567890123456";
        v_model = loadAndDecrypt_ONNX_ENNX_Model(input, key, iv);
    }
    catch (const std::runtime_error& e) {
        std::cerr << "Error loading model: " << e.what() << std::endl;
        return false;
    }

    ov::Core core;
    const std::string model_data(v_model.begin(), v_model.end());

    const auto [xml_, bin_] = coreModel_to_XB_buffer(core, model_data);

    std::vector<char> xml_buffer = xml_;
    std::vector<char> bin_buffer = bin_;

    info.engine_size = xml_buffer.size() + bin_buffer.size();
    info.gpu_model[0] = *"Intel core";

    const std::string encrypt_file = input.trans_path;
    encrypt_XML_BIN_to_hdats(
        xml_buffer,
        bin_buffer,
        info,
        encrypt_file);

    remove_temp_model_files();
    return true;
}
